
import torch
from torch import nn, Tensor
from torchvision.models import resnet152, resnet50, resnext101_32x8d


class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, x):
        shape = torch.prod(torch.tensor(x.shape[1:])).item()
        return x.view(-1, shape)



def makeNoClassiferResNet50():
    model = resnet50(pretrained=True)
    model = nn.Sequential(
        *list(model.children())[:-1],
        nn.BatchNorm2d(2048),
        Flatten(),
    )

    return model



def makeResNet50(classNum: int):
    model = nn.Sequential(
        makeNoClassiferResNet50(),
        nn.Linear(2048, classNum)
    )

    return model



def makeNoClassiferResNet152():
    model = resnet152(pretrained=True)
    model = nn.Sequential(
        *list(model.children())[:-1],
        # torch.nn.BatchNorm2d(2048),
        Flatten()
    )

    return model



def makeResNet152(classNum: int):
    model = nn.Sequential(
        makeNoClassiferResNet152(),
        nn.Linear(2048, classNum)
    )

    return model



def makeNoClassiferResNext101_32x8d():
    model = resnext101_32x8d(pretrained=True)
    model = nn.Sequential(
        *list(model.children())[:-1],
        torch.nn.BatchNorm2d(2048),
        Flatten()
    )

    return model



def makeResNext101_32x8d(classNum: int):
    model = nn.Sequential(
        makeNoClassiferResNet152(),
        nn.Linear(2048, classNum)
    )

    return model



class Model(nn.Module):
    def __init__(self) -> None:
        super(Model, self).__init__()
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=8, kernel_size=5),
            nn.MaxPool2d(kernel_size=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=8, out_channels=16, kernel_size=5),
            nn.MaxPool2d(kernel_size=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3),
            nn.MaxPool2d(kernel_size=2),
            nn.ReLU(),
        )

    def forward(self, inp: Tensor) -> Tensor:
        feature = self.feature_extractor(inp)
        print(feature.shape)
        feature = feature.mean(dim=[2, 3])
        print(feature.shape)
        return nn.functional.normalize(feature)


class CircleModel(nn.Module):
    def __init__(self) -> None:
        super(CircleModel, self).__init__()
        backBone = resnet152(pretrained=True)
        self.feature_extractor = nn.Sequential(
            *list(backBone.children())[:-1],
        )


    def forward(self, inp: Tensor) -> Tensor:
        feature = self.feature_extractor(inp)
        # print(feature.shape)
        feature = feature.mean(dim=[2, 3])
        # print(feature.shape)
        return nn.functional.normalize(feature)



if __name__ == '__main__':
    device = torch.device("cuda")
    # model = Model(300).to(device)
    model = CircleModel().to(device)
    print(model)

    x = torch.zeros([32, 3, 224, 224]).to(device)
    y = model(x)
    print(y.shape)

